Review updated until commit 62109c2

Description

• Implement AllToAll collective communication primitive with full backend integration

• Add lowering logic to convert AllToAll operations to communication primitives

• Include AllToAll in communication type handling and layout management

• Add comprehensive test coverage for multi-device AllToAll functionality

c10::intrusive_ptr<c10d::Work> postAllToAll(
    Communication* communication,
    DeviceIdxType my_device_index,
    c10d::Backend* backend,
    at::Tensor input_tensor,
    at::Tensor output_tensor) {
  NVF_ERROR(
      isTvContiguous(communication->in()),
      "Input tensor is not contiguous: ",
      communication->in(),
      " contiguity: ",
      communication->in()->domain()->getContiguityString());
  NVF_ERROR(
      isTvContiguous(communication->out()),
      "Output tensor is not contiguous: ",
      communication->out(),
      " contiguity: ",
      communication->out()->domain()->getContiguityString());

  // input_tv = [DIDx(d), n/d, m, ...]
  // output_tv = [n, DIDx(d), m/d, ...]
  // `n`: gathered dimension
  // `m`: scattered dimension
  // For alltoall correctness, we split `m` and reorder as [DIDx(d), d, n/d,
  // m/d, ...] such that alltoall_base splits across the `d` dimension.

  int64_t d = communication->team_size();
  auto input_sizes = input_tensor.sizes();

  NVF_CHECK(
      input_sizes.at(1) % d == 0,
      "Scattered dimension must be divisible by the team size");

  std::vector<int64_t> input_reshape_sizes(
      input_sizes.begin(), input_sizes.end());
  input_reshape_sizes.at(1) = d;
  input_reshape_sizes.insert(
      input_reshape_sizes.begin() + 2, input_sizes.at(1) / d);
  auto reshaped_input = input_tensor.reshape(input_reshape_sizes);

  std::vector<int64_t> permute_dims(input_reshape_sizes.size());
  std::iota(permute_dims.begin(), permute_dims.end(), 0);
  std::swap(permute_dims[0], permute_dims[1]);

  auto reordered_input = reshaped_input.permute(permute_dims).contiguous();

  auto flattened_input_tensor = viewAsCompact(reordered_input);
  auto flattened_output_tensor = viewAsCompact(output_tensor);

  // alltoall_base requires even splits of the input and output tensors.
  auto input_splits = at::tensor_split(
      flattened_input_tensor, communication->team_size(), /*dim=*/0);
  auto output_splits = at::tensor_split(
      flattened_output_tensor, communication->team_size(), /*dim=*/0);
  assertBuffersHaveSameSize(input_splits, output_splits);

  std::vector<int64_t> empty_split_sizes;
  return backend->alltoall_base(
      flattened_output_tensor,
      flattened_input_tensor,
      empty_split_sizes,
      empty_split_sizes,
      /*options=*/{});
}

if (c_logical_id == p2c_map.at(p_logical_id)) {
  fill_communication_info(
      CommunicationType::SendRecv, p_logical_id, c_logical_id);
} else {
  fill_communication_info(
      CommunicationType::AllToAll, nullptr, nullptr);
}

!test

Greptile Overview

Greptile Summary

This PR implements AllToAll collective communication support for nvfuser's multidevice framework. AllToAll enables simultaneous gathering along one dimension and scattering along another, effectively resharding tensors across devices.

Key Implementation Components

Communication Type Addition: Adds AllToAll enum value to CommunicationType and integrates it throughout the communication pipeline (operator overloads, type checking functions).

Core Runtime Implementation: The postAllToAll() function implements the NCCL alltoall_base primitive with tensor reshaping and permutation logic to organize data by destination rank. The implementation handles contiguous tensor validation and flattening.

Detection Logic: AllToAll is detected when both producer and consumer tensors are sharded but on different logical dimensions (c_logical_id != p2c_map[p_logical_id]). The detection sets both p_sharded_id and c_sharded_id to nullptr, which differs from other communication types.

Layout Handling: Includes an explicit TODO comment indicating that layout support for AllToAll needs fixes. Currently returns early in getCommunicationLayout() without enforcing specific layout requirements. The test uses permute operations as a workaround for allocation domain limitations.

Areas Requiring Attention

1. Documentation Inconsistency: Comments in postAllToAll() describing tensor shapes don't clearly distinguish between logical sharding representation and actual runtime memory layout, which could confuse future maintainers.

2. Nullptr Sharded IDs: Using nullptr for both sharded IDs in the AllToAll detection creates a fragile dependency where getCommunicationLayout() must explicitly handle AllToAll to avoid dereferencing null pointers.

3. Layout TODO: The acknowledged TODO for fixing AllToAll layout suggests incomplete implementation that may limit the operation's applicability or require workarounds like explicit permutations.

4. Limited Validation: The implementation validates that the scattered dimension is divisible by team size but doesn't validate output tensor dimensions or the relationship between gathered/scattered dimensions upfront.

Confidence Score: 3/5

• This PR is functionally correct for the tested use case but has implementation limitations and documentation issues that should be addressed
• The implementation passes its test case and follows existing patterns for other communication operations. However, several concerns reduce confidence: (1) the acknowledged TODO comment for layout fixes indicates incomplete implementation, (2) using nullptr for sharded_ids creates fragile coupling with getCommunicationLayout(), (3) inconsistent comments could mislead future maintainers, and (4) validation could be more comprehensive. The test coverage is adequate but only validates one scenario. These are not critical bugs but technical debt that limits robustness.
• Pay close attention to csrc/multidevice/communication.cpp (comment accuracy) and csrc/host_ir/lower_to_communication.cpp (nullptr handling and TODO resolution)

Important Files Changed

File Analysis

Sequence Diagram

sequenceDiagram
    participant User as User Code
    participant FD as FusionDefinition
    participant Lower as lower_to_communication
    participant Comm as Communication
    participant NCCL as NCCL Backend
    
    User->>FD: define AllToAll fusion with permute ops
    User->>FD: set sharding on different dims
    FD->>Lower: getCommunicationInfo(expr)
    Lower->>Lower: detect p_sharded && c_sharded
    Lower->>Lower: check if c_logical_id != p2c_map[p_logical_id]
    Lower->>Lower: return AllToAll with nullptr sharded_ids
    
    FD->>Lower: convertSingleOpToCommunication()
    Lower->>Lower: getCommunicationLayout() returns early for AllToAll
    Lower->>Lower: lowerToAllToAll()
    Lower->>Comm: create Communication(AllToAll)
    
    User->>FD: execute([input_tensor])
    FD->>Comm: postSingleCommunication()
    Comm->>Comm: postAllToAll()
    Comm->>Comm: check isTvContiguous for input/output
    Comm->>Comm: reshape input to split scattered dim
    Comm->>Comm: permute to move DIDx(d) outermost
    Comm->>Comm: flatten tensors with viewAsCompact
    Comm->>NCCL: alltoall_base(output, input)
    NCCL-->>Comm: Work handle
    Comm-->>User: output tensor with resharded data